# -*- coding: utf-8 -*-
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
代码生成任务收集器

负责收集所有需要生成的代码任务，并按照源码顺序排序。
"""

from dataclasses import dataclass
from typing import List, Optional

from ..ir import IRNode, IRModule
from ..ir.declarations.function import IRFunctionNode


@dataclass
class CodeGenTask:
    """代码生成任务"""

    node: IRNode
    source_line: int
    task_type: str  # "decl" 或 "impl"
    is_nested: bool
    parent_node: Optional[IRNode] = None  # 如果是嵌套函数，指向父函数
    priority: int = 0  # 优先级，用于调整顺序（嵌套函数优先级更高）


class CodeGenTaskCollector:
    """代码生成任务收集器"""

    def collect_tasks(self, ir_module: IRModule, for_source: bool) -> List[CodeGenTask]:
        """收集所有代码生成任务

        Args:
            ir_module: IR 模块
            for_source: 是否用于源文件

        Returns:
            排序后的任务列表
        """
        tasks = []

        # 遍历所有顶层节点
        for node in ir_module.nodes:
            # 收集该节点的任务
            node_tasks = self._collect_node_tasks(node, for_source, parent=None)
            tasks.extend(node_tasks)

        # 排序：确保嵌套函数总是排在父函数之前
        # 1. 先按父函数的行号排序（如果任务有父函数，使用父函数的行号；否则使用自己的行号）
        #    这样确保同一父函数下的所有任务（包括嵌套函数和父函数本身）会在一起
        # 2. 再按优先级排序（嵌套函数优先级更高，会排在前面）
        # 3. 最后按行号排序（确保同一父函数下的嵌套函数按行号排序）
        def get_sort_key(task: CodeGenTask) -> tuple:
            # 如果有父函数，使用父函数的行号作为主要排序键
            if task.parent_node and hasattr(task.parent_node, "source_line"):
                parent_line = getattr(task.parent_node, "source_line", 0) or 0
                return (parent_line, -task.priority, task.source_line)
            else:
                # 没有父函数，使用自己的行号
                return (task.source_line, -task.priority, task.source_line)

        tasks.sort(key=get_sort_key)

        return tasks

    def _collect_node_tasks(
        self, node: IRNode, for_source: bool, parent: Optional[IRNode]
    ) -> List[CodeGenTask]:
        """收集单个节点的任务（递归处理嵌套函数）

        Args:
            node: 节点
            for_source: 是否用于源文件
            parent: 父节点（如果是嵌套函数）

        Returns:
            任务列表
        """
        tasks = []

        # 如果是函数节点，先收集嵌套函数的任务
        if isinstance(node, IRFunctionNode):
            # 先收集嵌套函数的任务（优先级更高，会排在前面）
            nested_tasks = self._collect_nested_function_tasks(node, for_source)
            tasks.extend(nested_tasks)

            # 再收集当前函数的任务
            source_line = getattr(node, "source_line", 0) or 0
            task = CodeGenTask(
                node=node,
                source_line=source_line,
                task_type="impl" if for_source else "decl",
                is_nested=(parent is not None),
                parent_node=parent,
                priority=0,
            )
            tasks.append(task)
        else:
            # 其他类型的节点，直接创建任务
            source_line = getattr(node, "source_line", 0) or 0
            task = CodeGenTask(
                node=node,
                source_line=source_line,
                task_type="impl" if for_source else "decl",
                is_nested=False,
                parent_node=parent,
                priority=0,
            )
            tasks.append(task)

        return tasks

    def _collect_nested_function_tasks(
        self, parent_func: IRFunctionNode, for_source: bool
    ) -> List[CodeGenTask]:
        """收集函数内部嵌套函数的任务

        Args:
            parent_func: 父函数
            for_source: 是否用于源文件

        Returns:
            任务列表
        """
        tasks = []

        # 从函数体中收集嵌套函数
        nested_funcs = self._collect_nested_functions(parent_func)

        # 按源码行号排序
        nested_funcs.sort(key=lambda f: getattr(f, "source_line", 0) or 0)

        # 为每个嵌套函数创建任务
        for nested_func in nested_funcs:
            # 递归收集嵌套函数的嵌套函数
            nested_nested_tasks = self._collect_nested_function_tasks(
                nested_func, for_source
            )
            tasks.extend(nested_nested_tasks)

            # 创建嵌套函数的任务（优先级更高）
            source_line = getattr(nested_func, "source_line", 0) or 0
            task = CodeGenTask(
                node=nested_func,
                source_line=source_line,
                task_type="impl" if for_source else "decl",
                is_nested=True,
                parent_node=parent_func,
                priority=1,  # 嵌套函数优先级更高
            )
            tasks.append(task)

        return tasks

    def _collect_nested_functions(
        self, parent_func: IRFunctionNode
    ) -> List[IRFunctionNode]:
        """收集函数内部的嵌套函数

        Args:
            parent_func: 父函数

        Returns:
            嵌套函数列表
        """
        nested_funcs = []

        if not parent_func.body:
            return nested_funcs

        # 递归遍历函数体
        def collect_from_stmt(stmt):
            if isinstance(stmt, IRFunctionNode):
                # 检查是否是直接嵌套函数
                if self._is_direct_nested_of(stmt, parent_func):
                    nested_funcs.append(stmt)
            elif hasattr(stmt, "statements"):
                for child_stmt in stmt.statements:
                    collect_from_stmt(child_stmt)
            elif hasattr(stmt, "get_child_nodes"):
                for child in stmt.get_child_nodes():
                    collect_from_stmt(child)
            elif hasattr(stmt, "__dict__"):
                # 遍历所有属性
                for attr_value in stmt.__dict__.values():
                    if isinstance(attr_value, IRNode):
                        collect_from_stmt(attr_value)
                    elif isinstance(attr_value, list):
                        for item in attr_value:
                            if isinstance(item, IRNode):
                                collect_from_stmt(item)

        if hasattr(parent_func.body, "statements"):
            for stmt in parent_func.body.statements:
                collect_from_stmt(stmt)

        return nested_funcs

    def _is_direct_nested_of(
        self, nested_func: IRFunctionNode, parent_func: IRFunctionNode
    ) -> bool:
        """检查 nested_func 是否是 parent_func 的直接嵌套函数

        Args:
            nested_func: 嵌套函数
            parent_func: 父函数

        Returns:
            如果是直接嵌套函数则返回 True
        """
        parent = nested_func.parent
        while parent:
            if parent == parent_func:
                return True
            if isinstance(parent, IRFunctionNode):
                return False  # 是其他函数的嵌套函数
            parent = parent.parent
        return False
